from paddlenlp.datasets import load_dataset
from paddlenlp.transformers import BertTokenizer, BertModel, BertForTokenClassification
import paddle
from paddle.io import Dataset, DataLoader
import numpy as np
from PoemData import PoemData
import re
from PoetryGen import PoetryGen
from paddle.nn import Layer, Linear, Softmax
from paddle.static import InputSpec
from paddlenlp.metrics import Perplexity
from paddle.optimizer import AdamW



# 编写导入数据的map处理函数
def read(data_path):
    with open(data_path, 'r', encoding='utf-8') as f:
        # 跳过列名
        next(f)
        for line in f:
            yield line.replace("\n", "")


# 利用飞浆的load_dataset加载本地数据集
train_dataset = load_dataset(read, data_path='./datasets/train.txt', lazy=False)
test_dataset = load_dataset(read, data_path='./datasets/test.txt', lazy=False)
dev_dataset = load_dataset(read, data_path='./datasets/dev.txt', lazy=False)
# 打印
print('train_dataset 的样本数量：%d' % len(train_dataset))
print('单样本示例：%s' % train_dataset[0])


# 加载bert分词器-中文模型
bert_tokenizer = BertTokenizer.from_pretrained('bert-base-chinese')


# 处理效果展示
for poem in test_dataset[0:2]:
    token_poem, _ = bert_tokenizer.encode(poem).values()
    print(poem)
    print(token_poem)
    print(''.join(bert_tokenizer.convert_ids_to_tokens(token_poem)))


# 定义损失函数
class PoetryBertModelLossCriterion(Layer):
    def forward(self, pred_logits, label, input_mask):
        loss = paddle.nn.functional.cross_entropy(pred_logits, label, ignore_index=0, reduction='none')
        masked_loss = paddle.mean(loss * input_mask, axis=0)
        return paddle.sum(masked_loss)


# 创建网络
from Poetry_Bert_Model import PoetryBertModel
net = PoetryBertModel('bert-base-chinese', 128)
# 传入模型的数据集添加描述信息...
"""
InputSpec:
    用于描述模型输入的签名信息，包括 shape、dtype 和 name。
    此接口常用于指定高层 API 中模型的输入 Tensor 信息，或动态图转静态图时，指定被 paddle.jit.to_static 装饰的 forward 函数每个输入参数的 Tensor 信息。
"""

token_ids = InputSpec((-1, 128), 'int64', 'token')
token_type_ids = InputSpec((-1, 128), 'int64', 'token_type')
input_mask = InputSpec((-1, 128), 'float32', 'input_mask')
label = InputSpec((-1, 128), 'int64', 'label')

# 输入特征
inputs = [token_ids, token_type_ids, input_mask]
# 标签
labels = [label, input_mask]

'''
Model 对象是一个具备训练、测试、推理的神经网络。该对象同时支持静态图和动态图模式，飞桨框架默认为动态图模式，
通过 paddle.enable_static() 来切换到静态图模式。需要注意的是，需要在实例化 Model 对象之前完成切换。

参数
network (paddle.nn.Layer) - 是 paddle.nn.Layer 的一个实例。
inputs (InputSpec|list|tuple|dict|None，可选) - network 的输入，可以是 InputSpec 的实例，
    或者是一个 InputSpec 的 list，或者是格式为 {name: InputSpec} 的 dict，或者为 None。默认值为 None。
labels (InputSpec|list|tuple|None，可选) - network 的标签，可以是 InputSpec 的实例，
    或者是一个 InputSpec 的 list，或者为 None。 默认值为 None。
'''
model = paddle.Model(net, inputs, labels)
'''
配置模型所需的部件，比如优化器、损失函数和评价指标。

参数
optimizer (OOptimizer|None，可选) - 当训练模型的，该参数必须被设定。当评估或测试的时候，该参数可以不设定。默认值：None。
loss (Loss|Callable|None，可选) - 当训练模型的，该参数必须被设定。默认值：None。
metrics (Metric|list[Metric]|None，可选) - 当该参数被设定时，所有给定的评估方法会在训练和测试时被运行，并返回对应的指标。默认值：None。
amp_configs (str|dict|None，可选) - 混合精度训练的配置，通常是个 dict，也可以是 str。
    当使用自动混合精度训练或者纯 float16 训练时，amp_configs 的 key level 需要被设置为 O1 或者 O2，float32 训练时则默认为 O0。
    除了 level ，还可以传入更多的和混合精度 API 一致的参数，
    例如：init_loss_scaling、 incr_ratio 、 decr_ratio、 incr_every_n_steps、 decr_every_n_nan_or_inf、 use_dynamic_loss_scaling、 
        custom_white_list、 custom_black_list ，在静态图下还支持传入 custom_black_varnames 和 use_fp16_guard 。
        详细使用方法可以参考参考混合精度 API 的文档 auto_cast 和 GradScaler 。
        为了方便起见，当不设置其他的配置参数时，也可以直接传入 'O1' 、'O2' 。在使用 float32 训练时，该参数可以为 None。默认值：None。
   
优化器：AdamW     
    集成一阶动量和二阶动量的优化器Adam。AdamW其实是在Adam的基础上加入了weight decay正则化
'''

model.prepare(optimizer=AdamW(learning_rate=0.0001, parameters=model.parameters()),
              loss=PoetryBertModelLossCriterion(),
              metrics=[Perplexity()])

"""
Perplexity  困惑度是用来评价语言模型好坏的指标。语言模型是衡量句子好坏的模型，本质上是计算句子的概率。
     困惑度与测试集上的句子概率相关，其基本思想是：给测试集的句子赋予较高概率值的语言模型较好,当语言模型训练完之后，
        测试集中的句子都是正常的句子，那么训练好的模型就是在测试集上的概率越高越好
"""

"""
summary
    打印网络的基础结构和参数信息。
参数
input_size (tuple|InputSpec|list[tuple|InputSpec]，可选) - 输入 Tensor 的大小。如果网络只有一个输入，那么该值需要设定为 tuple 或 InputSpec。如果模型有多个输入。
    那么该值需要设定为 list[tuple|InputSpec]，包含每个输入的 shape 。如果该值没有设置，会将 self._inputs 作为输入。默认值：None。
dtype (str，可选) - 输入 Tensor 的数据类型，如果没有给定，默认使用 float32 类型。默认值：None。

返回
字典：包含网络全部参数的大小和全部可训练参数的大小。
"""
model.summary(inputs, [input.dtype for input in inputs])

# 加载数据集到数据集控制器分批次
train_loader = DataLoader(PoemData(train_dataset, bert_tokenizer, 128), batch_size=32, shuffle=True)
dev_loader = DataLoader(PoemData(dev_dataset, bert_tokenizer, 128), batch_size=32, shuffle=True)
model.fit(train_data=train_loader, epochs=10, save_dir='./checkpoint_epoch10', save_freq=1, verbose=1, eval_data=dev_loader, eval_freq=1)

# 保存网络结构，可在可视化界面查看
"""
jit.save
    将输入的 Layer 或 function 存储为 paddle.jit.TranslatedLayer 格式的模型，载入后可用于预测推理或者 fine-tune 训练。
    paddle.Model.save的第一个参数需要设置为待保存的模型和参数等文件的前缀名，第二个参数 training 表示是否保存动态图模型以继续训练，
        默认是 True，这里需要设为 False，即保存推理部署所需的参数与文件。接前文高层 API 训练的示例代码，保存推理模型代码示例如下：
"""
# paddle.jit.save(net, './runs_False/model', inputs)


# 开启模型训练
"""
fit
    训练模型。当 eval_data 给定时，会在 eval_freq 个 epoch 后进行一次评估。
参数
train_data (Dataset|DataLoader，可选) - 一个可迭代的数据源，推荐给定一个 paddle paddle.io.Dataset 或 paddle.io.Dataloader 的实例。默认值：None。
eval_data (Dataset|DataLoader，可选) - 一个可迭代的数据源，推荐给定一个 paddle paddle.io.Dataset 或 paddle.io.Dataloader 的实例。
    当给定时，会在每个 epoch 后都会进行评估。默认值：None。
batch_size (int，可选) - 训练数据或评估数据的批大小，当 train_data 或 eval_data 为 DataLoader 的实例时，该参数会被忽略。默认值：1。
epochs (int，可选) - 训练的轮数。默认值：1。
eval_freq (int，可选) - 评估的频率，多少个 epoch 评估一次。默认值：1。
log_freq (int，可选) - 日志打印的频率，多少个 step 打印一次日志。默认值：10。
save_dir (str|None，可选) - 保存模型的文件夹，如果不设定，将不保存模型。默认值：None。
save_freq (int，可选) - 保存模型的频率，多少个 epoch 保存一次模型。默认值：1。
verbose (int，可选) - 可视化的模型，必须为 0，1，2。当设定为 0 时，不打印日志，设定为 1 时，使用进度条的方式打印日志，设定为 2 时，一行一行地打印日志。默认值：2。
drop_last (bool，可选) - 是否丢弃训练数据中最后几个不足设定的批次大小的数据。默认值：False。
shuffle (bool，可选) - 是否对训练数据进行洗牌。当 train_data 为 DataLoader 的实例时，该参数会被忽略。默认值：True。
num_workers (int，可选) - 启动子进程用于读取数据的数量。当 train_data 和 eval_data 都为 DataLoader 的实例时，该参数会被忽略。默认值：0。
callbacks (Callback|list[Callback]|None，可选) - Callback 的一个实例或实例列表。该参数不给定时，
    默认会插入 ProgBarLogger 和 ModelCheckpoint 这两个实例。默认值：None。
accumulate_grad_batches (int，可选) - 训练过程中，在优化器更新之前累积梯度的批次数。通过改变该参数可以模拟大的 batch_size。默认值：1。
num_iters (int，可选) - 训练模型过程中的迭代次数。如果设置为 None，则根据参数 epochs 来训练模型，否则训练模型 num_iters 次。默认值：None。
"""
# model.fit(train_data=train_loader,
#           eval_data=dev_loader,
#           epochs=2,
#           save_dir='checkpoint_2',
#           save_freq=1,
#           eval_freq=1,
#           verbose=1)


"""
    模型预测
        功能：
            1.  随机生成一首诗
            2.  藏头诗
            3.  诗词续写
            4.  生成特定风格的诗
"""
# 导入模型准备进行预测
net = PoetryBertModel('bert-base-chinese', 128)
model = paddle.Model(net)
model.load('./checkpoint_2/final.pdopt')
poetry_gen = PoetryGen(model, bert_tokenizer)


# 遇到标点符号换行的工具函数
def poetry_show(poetry):
    pattern = r"([，。；？])"
    text = re.sub(pattern, r'\1 ', poetry)
    for p in text.split():
        if p:
            print(p)


# 随机生成一首诗
poetry = poetry_gen.generate()
poetry_show(poetry)

# 生成藏头诗
poetry = poetry_gen.generate(head=['这', '个', '魔', '型', '很', '无', '敌'])
poetry_show(poetry)

# 生成特定风格的诗
poetry = poetry_gen.generate(style='会当凌绝顶，一览众山小。')
poetry_show(poetry)

# 生成特定开头的诗
poetry = poetry_gen.generate(head='好好学习')
poetry_show(poetry)
